#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2020/7/2 13:54
# @USER    : Shengji He
# @File    : KthSmallestElementInSortedMatrix.py
# @Software: PyCharm
# @Version  : Python-
# @TASK:
from typing import List
import heapq


class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        """
        Given a n x n matrix where each of the rows and columns are sorted in ascending order, find the kth
        smallest element in the matrix.

        Note that it is the kth smallest element in the sorted order, not the kth distinct element.

        Example:
            matrix =
                [
                    [ 1,  5,  9],

                    [10, 11, 13],

                    [12, 13, 15]
                ],

            k = 8,

            return 13.
        Note:
            - You may assume k is always valid, 1 ≤ k ≤ n2.

        :param matrix:
        :param k:
        :return:
        """
        # def helper(coord):
        #     return matrix[coord[0]][coord[1]]
        if not matrix or not matrix[0]:
            return 0
        # candidate_list = []
        # rows, columns = len(matrix), len(matrix[0])
        # dp = []
        # for i in range(1, k + 1):
        #     if i == 1:
        #         dp.append((0, 0))
        #         if 1 < rows:
        #             candidate_list.append((1, 0))
        #         if 1 < columns:
        #             candidate_list.append((0, 1))
        #         continue
        #     candidate_value = list(map(helper, candidate_list))
        #     index = candidate_value.index(min(candidate_value))
        #     dp.append(candidate_list[index])
        #     if candidate_list[index][0] + 1 < rows and \
        #             (candidate_list[index][0] + 1, candidate_list[index][1]) not in candidate_list:
        #         candidate_list.append((candidate_list[index][0] + 1, candidate_list[index][1]))
        #     if candidate_list[index][1] + 1 < columns and \
        #             (candidate_list[index][0], candidate_list[index][1] + 1) not in candidate_list:
        #         candidate_list.append((candidate_list[index][0], candidate_list[index][1] + 1))
        #     candidate_list.pop(index)
        # return matrix[dp[-1][0]][dp[-1][1]]

        # way 2
        n = len(matrix)
        pq = [(matrix[i][0], i, 0) for i in range(n)]
        heapq.heapify(pq)
        for i in range(k - 1):
            num, x, y = heapq.heappop(pq)
            if y != n - 1:
                heapq.heappush(pq, (matrix[x][y + 1], x, y + 1))
        return heapq.heappop(pq)[0]


if __name__ == '__main__':
    S = Solution()
    matrix = [
        [1, 5, 9],
        [10, 11, 13],
        [12, 13, 15],
    ]
    k = 8
    print(S.kthSmallest(matrix, k))
    print('done')
